"""
    DiagnosticWriters

Write GSI-compatible diagnostic output files.

This module provides functions to write diagnostic files in Fortran GSI format:
- Observation space diagnostics (O-B, O-A)
- Analysis increment files
- Fit statistics (fort.* files)
- Innovation diagnostics
- QC summary files

All output is binary-compatible with Fortran GSI for validation purposes.
"""
module DiagnosticWriters

using Printf
using Dates
using Statistics

# Forward declarations - modules must be loaded before this one
# Do not include them here to avoid type conflicts

export write_diag_obs_anl, write_diag_obs_ges
export write_innovation_diagnostics_binary
export write_fit_statistics_binary
export write_qc_diagnostics_binary
export write_all_diagnostics
export DiagnosticOutput

"""
    DiagnosticOutput

Container for all diagnostic information during analysis.
"""
mutable struct DiagnosticOutput
    output_dir::String
    analysis_time::DateTime
    grid_size::Tuple{Int, Int, Int}

    # Observation data
    observations::Dict{String, Vector{ObservationRecord}}

    # Statistics
    innovation_stats::Dict{String, Dict{String, Float64}}
    qc_stats::Dict{String, Dict{String, Int}}

    # Convergence info
    cost_history::Vector{Float64}
    gradient_norms::Vector{Float64}
    iterations::Int
    converged::Bool

    # Analysis increments
    increment_fields::Dict{String, Array{Float64}}
end

"""
    DiagnosticOutput(output_dir::String, analysis_time::DateTime, grid_size::Tuple)

Create new diagnostic output container.
"""
function DiagnosticOutput(output_dir::String, analysis_time::DateTime, grid_size::Tuple)
    mkpath(output_dir)

    return DiagnosticOutput(
        output_dir,
        analysis_time,
        grid_size,
        Dict{String, Vector{ObservationRecord}}(),
        Dict{String, Dict{String, Float64}}(),
        Dict{String, Dict{String, Int}}(),
        Float64[],
        Float64[],
        0,
        false,
        Dict{String, Array{Float64}}()
    )
end

"""
    write_diag_obs_anl(filename::String, observations::Vector{ObservationRecord},
                       grid_meta::GridMetadata, analysis_time::DateTime)

Write analysis observation diagnostics file (diag_obs_setup_anl.YYYYMMDDHH).

This file contains observation-space diagnostics after analysis:
- O-A (observation minus analysis) for each observation
- Analysis values at observation locations
- QC flags and weights
"""
function write_diag_obs_anl(filename::String, observations::Vector{ObservationRecord},
                            grid_meta::GridMetadata, analysis_time::DateTime)

    println("Writing analysis observation diagnostics: $filename")

    # Determine variable name from observations
    var_name = length(observations) > 0 ? "ps" : "ps"  # Default to ps

    ff = open_fortran_file(filename, "w", byte_order=:BigEndian)

    try
        # Write header
        write_fortran_header(ff, var_name, var_name)

        # Write grid metadata
        write_grid_metadata(ff, grid_meta)

        # Group observations by type
        obs_by_type = group_observations_by_type(observations)

        # Write observation records for each type
        for (obs_type, obs_list) in obs_by_type
            write_observation_group(ff, obs_list, obs_type)
        end

        println("  Wrote $(length(observations)) observations")
        println("  Observation types: $(length(obs_by_type))")

    finally
        close_fortran_file(ff)
    end
end

"""
    write_diag_obs_ges(filename::String, observations::Vector{ObservationRecord},
                       grid_meta::GridMetadata)

Write background/guess observation diagnostics file (diag_obs_setup_ges.ensmean).

This file contains observation-space diagnostics before analysis:
- O-B (observation minus background) for each observation
- Background values at observation locations
- QC flags
"""
function write_diag_obs_ges(filename::String, observations::Vector{ObservationRecord},
                            grid_meta::GridMetadata)

    println("Writing background observation diagnostics: $filename")

    var_name = length(observations) > 0 ? "ps" : "ps"

    ff = open_fortran_file(filename, "w", byte_order=:BigEndian)

    try
        # Write header
        write_fortran_header(ff, var_name, var_name)

        # Write grid metadata
        write_grid_metadata(ff, grid_meta)

        # Group observations by type
        obs_by_type = group_observations_by_type(observations)

        # Write observation records for each type
        for (obs_type, obs_list) in obs_by_type
            write_observation_group(ff, obs_list, obs_type)
        end

        println("  Wrote $(length(observations)) observations")

    finally
        close_fortran_file(ff)
    end
end

"""
    write_grid_metadata(ff::FortranFile, grid_meta::GridMetadata)

Write grid metadata record to diagnostic file.
"""
function write_grid_metadata(ff::FortranFile, grid_meta::GridMetadata)
    # Pack grid metadata into byte array
    io = IOBuffer()

    write(io, hton(grid_meta.grid_type))
    write(io, hton(grid_meta.nx))
    write(io, hton(grid_meta.ny))
    write(io, hton(grid_meta.nz))
    write(io, hton(grid_meta.dx))
    write(io, hton(grid_meta.dy))
    write(io, hton(grid_meta.lon_min))
    write(io, hton(grid_meta.lat_min))

    # Write vertical coordinates
    for coord in grid_meta.vertical_coords
        write(io, hton(coord))
    end

    data = take!(io)
    write_fortran_record(ff, data)
end

"""
    write_observation_group(ff::FortranFile, observations::Vector{ObservationRecord},
                           obs_type::String)

Write a group of observations of the same type.
"""
function write_observation_group(ff::FortranFile, observations::Vector{ObservationRecord},
                                obs_type::String)

    # Pack all observations in this group into a single record
    io = IOBuffer()

    for obs in observations
        obs_bytes = pack_observation_record(obs)
        write(io, obs_bytes)
    end

    data = take!(io)

    if length(data) > 0
        write_fortran_record(ff, data)
    end
end

"""
    group_observations_by_type(observations::Vector{ObservationRecord})

Group observations by type code.
"""
function group_observations_by_type(observations::Vector{ObservationRecord})
    obs_groups = Dict{String, Vector{ObservationRecord}}()

    for obs in observations
        type_name = "type_$(obs.obs_type)"
        if !haskey(obs_groups, type_name)
            obs_groups[type_name] = ObservationRecord[]
        end
        push!(obs_groups[type_name], obs)
    end

    return obs_groups
end

"""
    write_fit_statistics_binary(filename::String, var_name::String,
                               stats::Dict{String, Float64}, iteration::Int=1)

Write fit statistics file in GSI fort.* format.

# Arguments
- `filename`: Output file path (e.g., "fort.201")
- `var_name`: Variable name ("ps", "t", "q", "uv")
- `stats`: Statistics dictionary with keys: count, bias, rms, cpen, qcpen
- `iteration`: Iteration number (default 1)
"""
function write_fit_statistics_binary(filename::String, var_name::String,
                                    stats::Dict{String, Float64}, iteration::Int=1)

    println("Writing fit statistics: $filename")

    open(filename, "w") do f
        # Write header
        if var_name == "ps"
            write(f, "current fit of surface pressure data, ranges in mb\n")
        elseif var_name == "t"
            write(f, "current fit of temperature data\n")
        elseif var_name == "q"
            write(f, "current fit of moisture data\n")
        elseif var_name == "uv"
            write(f, "current fit of wind data\n")
        else
            write(f, "current fit of $var_name data\n")
        end

        write(f, "-"^50 * "\n")

        # Write pressure range for ps
        if var_name == "ps"
            write(f, " pressure levels (hPa)=   0.0 2000.0\n")
        end

        write(f, " o-g it     obs use typ styp     count      bias       rms      cpen     qcpen\n")

        # Format output line
        count = Int(get(stats, "count", 0))
        bias = get(stats, "bias", 0.0)
        rms = get(stats, "rms", 0.0)
        cpen = get(stats, "cpen", 0.0)
        qcpen = get(stats, "qcpen", 0.0)

        @printf(f, " o-g %02d  %6s asm all      %6d  %8.4f  %8.4f  %8.4f  %8.4f\n",
                iteration, var_name, count, bias, rms, cpen, qcpen)

        # Write summary statistics
        n_rejected = Int(get(stats, "n_rejected", 0))
        write(f, " number of  $(var_name) obs that failed gross test = $(n_rejected) ")
        write(f, "nonlin qc test =     0\n")

        n_read = Int(get(stats, "n_read", count))
        n_keep = Int(get(stats, "n_keep", count))
        @printf(f, " type    %-4s jiter   %d nread  %8d nkeep %8d num %8d\n",
                var_name, iteration, n_read, n_keep, count)

        # Penalty statistics
        pen = cpen * count
        @printf(f, " type    %-4s pen=  0.%15.15fE+04 qcpen=  0.%15.15fE+04 r=  %.6f  qcr=  %.6f\n",
                var_name, pen/1e4, pen/1e4, rms, rms)
    end

    println("  → $filename")
end

"""
    write_innovation_diagnostics_binary(filename::String,
                                       observations::Vector{ObservationRecord})

Write innovation (O-B) diagnostics in text format.

This provides a human-readable summary of innovation statistics.
"""
function write_innovation_diagnostics_binary(filename::String,
                                            observations::Vector{ObservationRecord})

    println("Writing innovation diagnostics: $filename")

    # Group by type and compute statistics
    obs_by_type = group_observations_by_type(observations)

    open(filename, "w") do f
        write(f, "="^80 * "\n")
        write(f, "GSI INNOVATION DIAGNOSTICS (O-B)\n")
        write(f, "="^80 * "\n")
        write(f, "Generated: $(Dates.format(now(), "yyyy-mm-dd HH:MM:SS"))\n\n")

        write(f, @sprintf("%-20s %8s %10s %10s %10s %10s\n",
                         "Obs Type", "Count", "Mean", "RMS", "Bias", "StdDev"))
        write(f, "-"^80 * "\n")

        total_count = 0
        for (obs_type, obs_list) in sort(collect(obs_by_type))
            stats = compute_innovation_statistics(obs_list)

            @printf(f, "%-20s %8d %10.4f %10.4f %10.4f %10.4f\n",
                    obs_type,
                    stats["n_used"],
                    stats["mean_innovation"],
                    stats["rms_innovation"],
                    stats["bias"],
                    stats["std_innovation"])

            total_count += stats["n_used"]
        end

        write(f, "-"^80 * "\n")
        write(f, "Total observations used: $total_count\n")
        write(f, "="^80 * "\n")
    end

    println("  → $filename")
end

"""
    write_qc_diagnostics_binary(filename::String, qc_stats::Dict{String, Dict{String, Int}})

Write quality control diagnostics summary.
"""
function write_qc_diagnostics_binary(filename::String, qc_stats::Dict{String, Dict{String, Int}})

    println("Writing QC diagnostics: $filename")

    open(filename, "w") do f
        write(f, "="^80 * "\n")
        write(f, "GSI QUALITY CONTROL DIAGNOSTICS\n")
        write(f, "="^80 * "\n")
        write(f, "Generated: $(Dates.format(now(), "yyyy-mm-dd HH:MM:SS"))\n\n")

        write(f, @sprintf("%-20s %10s %10s %10s %10s\n",
                         "Obs Type", "Total", "Used", "Rejected", "Reject %"))
        write(f, "-"^80 * "\n")

        for (obs_type, stats) in sort(collect(qc_stats))
            total = get(stats, "total", 0)
            used = get(stats, "used", 0)
            rejected = get(stats, "rejected", 0)
            reject_pct = total > 0 ? 100.0 * rejected / total : 0.0

            @printf(f, "%-20s %10d %10d %10d %9.2f%%\n",
                    obs_type, total, used, rejected, reject_pct)
        end

        write(f, "="^80 * "\n")
    end

    println("  → $filename")
end

"""
    write_all_diagnostics(diag::DiagnosticOutput)

Write all diagnostic files for an analysis run.

This is the main entry point for generating complete GSI-compatible
diagnostic output.
"""
function write_all_diagnostics(diag::DiagnosticOutput)

    println("\n" * "="^80)
    println("WRITING GSI-COMPATIBLE DIAGNOSTIC FILES")
    println("="^80)

    # Create grid metadata
    grid_meta = GridMetadata(
        Int32(1),  # Regional grid
        Int32(diag.grid_size[1]),
        Int32(diag.grid_size[2]),
        Int32(diag.grid_size[3]),
        1.0,   # dx
        1.0,   # dy
        0.0,   # lon_min
        0.0,   # lat_min
        collect(1.0:Float64(diag.grid_size[3]))  # vertical coords
    )

    # Write observation diagnostics
    all_obs = ObservationRecord[]
    for (obs_type, obs_list) in diag.observations
        append!(all_obs, obs_list)
    end

    if !isempty(all_obs)
        # Analysis diagnostics
        analysis_time_str = Dates.format(diag.analysis_time, "yyyymmddHH")
        anl_file = joinpath(diag.output_dir, "diag_obs_setup_anl.$analysis_time_str")
        write_diag_obs_anl(anl_file, all_obs, grid_meta, diag.analysis_time)

        # Background diagnostics
        ges_file = joinpath(diag.output_dir, "diag_obs_setup_ges.ensmean")
        write_diag_obs_ges(ges_file, all_obs, grid_meta)

        # Innovation diagnostics (text)
        innov_file = joinpath(diag.output_dir, "innovation_diagnostics.txt")
        write_innovation_diagnostics_binary(innov_file, all_obs)

        # QC diagnostics
        if !isempty(diag.qc_stats)
            qc_file = joinpath(diag.output_dir, "qc_diagnostics.txt")
            write_qc_diagnostics_binary(qc_file, diag.qc_stats)
        end
    end

    # Write fit statistics (fort.* files)
    write_fit_statistics_files(diag)

    # Write convergence diagnostics
    write_convergence_file(diag)

    println("="^80)
    println("✓ All diagnostic files written to: $(diag.output_dir)")
    println("="^80)
end

"""
    write_fit_statistics_files(diag::DiagnosticOutput)

Write fit statistics fort.* files for all variables.
"""
function write_fit_statistics_files(diag::DiagnosticOutput)

    # Standard GSI fort file numbers for each variable
    fort_files = Dict(
        "ps" => 201,
        "uv" => 202,
        "t" => 203,
        "q" => 204
    )

    for (var_name, fort_num) in fort_files
        filename = joinpath(diag.output_dir, "fort.$fort_num")

        # Get statistics for this variable
        stats = get(diag.innovation_stats, var_name, Dict{String, Float64}())

        # Add default values if missing
        default_stats = Dict{String, Float64}(
            "count" => 0.0,
            "bias" => 0.0,
            "rms" => 0.0,
            "cpen" => 0.0,
            "qcpen" => 0.0,
            "n_rejected" => 0.0,
            "n_read" => 0.0,
            "n_keep" => 0.0
        )

        merged_stats = merge(default_stats, stats)

        write_fit_statistics_binary(filename, var_name, merged_stats, 1)
    end
end

"""
    write_convergence_file(diag::DiagnosticOutput)

Write convergence diagnostics file.
"""
function write_convergence_file(diag::DiagnosticOutput)

    filename = joinpath(diag.output_dir, "convergence_diagnostics.txt")

    println("Writing convergence diagnostics: $filename")

    open(filename, "w") do f
        write(f, "="^80 * "\n")
        write(f, "GSI CONVERGENCE DIAGNOSTICS\n")
        write(f, "="^80 * "\n")
        write(f, "Analysis time: $(diag.analysis_time)\n")
        write(f, "Generated: $(Dates.format(now(), "yyyy-mm-dd HH:MM:SS"))\n\n")

        write(f, "Iterations: $(diag.iterations)\n")
        write(f, "Converged: $(diag.converged)\n\n")

        if !isempty(diag.cost_history)
            initial_cost = diag.cost_history[1]
            final_cost = diag.cost_history[end]
            reduction = (initial_cost - final_cost) / initial_cost * 100

            @printf(f, "Initial cost function: %.6e\n", initial_cost)
            @printf(f, "Final cost function:   %.6e\n", final_cost)
            @printf(f, "Cost reduction:        %.2f%%\n\n", reduction)

            write(f, @sprintf("%10s %16s %16s\n", "Iteration", "Cost Function", "Gradient Norm"))
            write(f, "-"^80 * "\n")

            for i in 1:length(diag.cost_history)
                grad_norm = i <= length(diag.gradient_norms) ? diag.gradient_norms[i] : 0.0
                @printf(f, "%10d %16.6e %16.6e\n", i-1, diag.cost_history[i], grad_norm)
            end
        end

        write(f, "="^80 * "\n")
    end

    println("  → $filename")
end

end # module DiagnosticWriters
